import os
import torch
from torchvision.io import read_image
from torchvision.transforms import functional as F
from torch.utils.data import Dataset, DataLoader, random_split
from torchvision.transforms import Compose, Resize, Normalize
import random
import torchvision.transforms.functional as TF
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
import matplotlib.pyplot as plt
import numpy as np
import cv2


# 自定义数据集类
class RoadDataset(Dataset):
    def __init__(self, img_dir, anno_dir, transforms=None):
        # 图像文件夹路径
        self.img_dir = img_dir
        # 标注文件夹路径
        self.anno_dir = anno_dir
        # 数据变换
        self.transforms = transforms
        # 获取图像文件名列表
        self.img_files = os.listdir(img_dir)

    def __len__(self):
        # 返回数据集样本数量
        return len(self.img_files)

    def __getitem__(self, idx):
        # 构建图像文件路径
        img_path = os.path.join(self.img_dir, self.img_files[idx])
        # 读取图像
        image = read_image(img_path)

        # 构建标注文件路径，假设标注文件名与图像文件名对应，只是后缀不同
        anno_path = os.path.join(self.anno_dir, self.img_files[idx].replace('.png', '_L.png'))
        # 以灰度模式读取标注图像
        annotation = cv2.imread(anno_path, cv2.IMREAD_GRAYSCALE)
        # 将标注转换为torch的long类型张量
        annotation = torch.from_numpy(annotation).long()

        if self.transforms is not None:
            # 应用数据变换
            image, annotation = self.transforms(image, annotation)

        return image, annotation


# 自定义ToTensor变换类
class ToTensor:
    def __call__(self, image, annotation):
        # 将图像转换为张量
        image = F.to_tensor(image)
        # 将标注转换为张量
        annotation = torch.from_numpy(np.array(annotation)).long()
        return image, annotation


# 数据增强：随机水平翻转
class RandomHorizontalFlip:
    def __call__(self, image, annotation):
        if random.random() < 0.5:
            # 水平翻转图像
            image = TF.hflip(image)
            # 水平翻转标注
            annotation = TF.hflip(annotation)
        return image, annotation


# ****构建U-Net模型的DoubleConv模块，前面的模型都是直接套用别人的成品****
class DoubleConv(nn.Module):#nn.Moudle是一个DL中常用于构建卷积神经网络CNN的模块
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        # 连续两个卷积层、批归一化层和ReLU激活函数
        """
            nn.Sequential对象,是一个有序的模块容器,按顺序包装多个神经网络层。
            nn.Sequential会自动管理这些层向前传播。将输入数据传递给model时,数据会按照输入数据依次通过这些层
            
            归一化原理：
            nn.BatchNorm2d是针对二维数据(如图像，[batch_size,channels,height,width])的批归一化层。
            每个小批量数据在每个通道维度上归一化，使得数据稳定，具体为计算均值方差然后正态分布归一，防止梯度爆炸。
        """
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),#定义第一个二维卷积层
            nn.BatchNorm2d(out_channels),#定义第一个批归一化层。
            nn.ReLU(inplace=True),#定义第一个ReLU激活函数，原地操作。ReLU函数将所有负数输入置为0，正数输入不变，为神经网络引入非线性
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        """
        定义前向传播函数，nn.Moudle类的重要方法，用于定义模型的前向传播逻辑。
        输入数据传递给模型实例时，实际上调用的就是这个forward函数。
        forward()函数中，需要描述输入数据如何通过模型的各个层，以及层与层之间的数据流动和处理方式。
        它决定了模型从输入到输出的计算过程，核心。
        """
        return self.conv(x)#将输入数据x传入之前定义nn.Sequential对象self.conv中，依次经过其中的各个模块（两个卷积层，两个批归一化层，两个ReLU激活函数），并返回最终的输出结果


# 构建U-Net模型的下采样模块
class Down(nn.Module):
    """
    下采样过程中通道数增加，因为下采样进行时图像空间尺寸减小，为了保留足够信息并捕捉不同类型的特征需要增加通道数。
    通常伴随卷积操作，kernel在不同通道上滑动并进行加权求和提取局部特征。增加通道数可以使得卷积在更多的特征维度进行，从而扩大感受野
    """
    def __init__(self, in_channels, out_channels):
        super(Down, self).__init__()
        # 最大池化后接DoubleConv
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),#尺寸减半
            DoubleConv(in_channels, out_channels)#卷积归一激活×2
        )

    def forward(self, x):
        return self.maxpool_conv(x)
"""
关于通道

1. 数据维度与信息载体
    物理层面:RGB
    抽象层面:通道是数据维度的一种体现，用于承载不同类型的信息。
            比如:在医学图像中，通道可能表示不同的成像模态;在一些经过预处理的图像数据中,通道可能表示不同的特征图，每个特征图是通过特定kernel提取的图像局部特征的响应
2. 特征提取与表示
    卷积操作与特征学习:在CNN中,卷积层通过卷积核对输入数据进行卷积操作提取特征。每个卷积核在输入的每个通道上滑动，并进行加权求和，生成一个新的特征图。
3. 通道与特征
    例如RGB有3个通道,kernel在这些通道上进行卷积提取出特征,这些新的特征会体现在输出的特征图中,which 同样具有特征维度
"""

# 构建U-Net模型的上采样模块
# 将经果下采样后缩小的特征图恢复到与输入图像相近的尺寸，以便逐像素分类。
# 恢复空间尺寸后，把下采样中不同阶段的特征（含细节信息）与上采样后的特征进行融合
class Up(nn.Module):
    def __init__(self, in_channels, out_channels, bilinear=True):
        """
        in_channels:输入特征图的通道数。通常是下采样通道数的两倍
        out_channels:经过该上采样模块处理后输出特征图的通道数
        bilinear:布尔值,用于选择上采样的方式。如果为True,使用双线性插值;如果为False,使用转置卷积
        """
        super(Up, self).__init__()
        if bilinear:
            # 双线性插值上采样
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        else:
            # 转置卷积(反卷积)上采样，4×4特征图-->2×2卷积核-->生成8×8特征图
            self.up = nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)

        # 双卷积模块
        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):#参数x1:下采样后图;参数x2:下采样前图
        x1 = self.up(x1)
        # 计算特征图尺寸差异，原始格式(batch_size,channels,height,width);
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        # 对x1进行填充
        x1 = nn.functional.pad(x1, [diffX // 2, diffX - diffX // 2,
                                    diffY // 2, diffY - diffY // 2])
        # 拼接特征图
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


# 构建U-Net模型的输出模块
class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        """ 1x1卷积层
        调整通道数:经过下采样、上采样、特征融合等操作后,得到特征图通道数可能与最终任务所需的通道数不一致。比如要10个而生成了通道为64的特征图,可调整出10个;反之也是。
        实现跨通道信息整合:虽然1×1卷积核只作用于单个像素位置,但它在所有通道上进行操作。这意味着它可以对每个像素位置的不同通道信息加权求和,实现跨通道的信息整合。
                         从数学角度看,对于输入特征图的每个像素位置,1×1卷积核将该位置在所有输入通道上的值进行线性组合,生成输出特征图对应像素位置在新通道上的值。
                            本质上,该卷积核是一个1×1×Cin的小型张量（Cin是输入特征图的通道数）
                            在进行卷积运算时，对于输入特征图的每个像素位置，1×1 卷积核会与该位置在所有输入通道上的值进行点乘，并将结果累加，再加上偏置项（如果有），得到输出特征图对应位置的值。
                            例如，在图像分类任务中，之前的卷积层可能在不同通道上分别提取了颜色、纹理、形状等特征，1×1 卷积核可以将这些特征综合起来，形成更具代表性的特征表示，有助于模型做出更准确的分类决策。
        """
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)


# 构建完整的U-Net模型
class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        # 输入卷积模块
        self.inc = DoubleConv(n_channels, 64)
        # 下采样模块
        self.down1 = Down(64, 128)#size÷2,扩大通道数
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        factor = 2 if bilinear else 1
        self.down4 = Down(512, 1024 // factor)#size÷2,双线性通道数不变，方便操作
        # 上采样模块
        self.up1 = Up(1024, 512 // factor, bilinear)
        self.up2 = Up(512, 256 // factor, bilinear)
        self.up3 = Up(256, 128 // factor, bilinear)
        self.up4 = Up(128, 64, bilinear)
        # 输出卷积模块
        self.outc = OutConv(64, n_classes)#一位卷积核整合

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits


# 数据预处理变换
transform = Compose([
    Resize((256, 256)),
    ToTensor(),
    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 数据增强变换
transform_with_aug = Compose([
    Resize((256, 256)),
    RandomHorizontalFlip(),
    ToTensor(),
    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 数据集路径
img_dir = 'data/train/images'
anno_dir = 'data/train/labels'

# 实例化数据集
dataset_with_aug = RoadDataset(img_dir, anno_dir, transforms=transform_with_aug)

# 划分数据集
train_size = int(0.7 * len(dataset_with_aug))
val_size = int(0.2 * len(dataset_with_aug))
test_size = len(dataset_with_aug) - train_size - val_size
train_dataset, val_dataset, test_dataset = random_split(dataset_with_aug, [train_size, val_size, test_size])

# 创建数据加载器
batch_size = 4
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 初始化模型
n_channels = 3  # RGB图像通道数
n_classes = 2  # 道路和非道路两类
model = UNet(n_channels, n_classes)

# 定义优化器
lr = 0.0001
optimizer = optim.Adam(model.parameters(), lr=lr)

# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 训练模型
num_epochs = 50
for epoch in range(num_epochs):
    running_loss = 0.0
    model.train()
    pbar = tqdm(train_loader)
    for i, (images, labels) in enumerate(pbar):
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()

        outputs = model(images)
        loss = criterion(outputs, labels)

        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        pbar.set_description(f'Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / (i + 1):.4f}')


# 在验证集上评估
model.eval()
total_iou = 0.0#交并比
total_pixel_acc = 0.0#像素准确率
num_batches = 0

with torch.no_grad():
    pbar = tqdm(val_loader)
    for i, (images, labels) in enumerate(pbar):
        images, labels = images.to(device), labels.to(device)

        outputs = model(images)
        _, predicted = torch.max(outputs, 1)

        intersection = (predicted == labels).sum().item()#布尔张量sum()求和，然后通过item()取出取标量
        union = (predicted >= 0).sum().item() + (labels >= 0).sum().item() - intersection
        batch_iou = intersection / union if union > 0 else 0

        total_iou += batch_iou
        total_pixel_acc += (predicted == labels).sum().item() / (labels.numel())

        num_batches += 1

mean_iou = total_iou / num_batches
mean_pixel_acc = total_pixel_acc / num_batches

print(f'Validation Mean IoU: {mean_iou:.4f}, Validation Pixel Accuracy: {mean_pixel_acc:.4f}')


# 可视化评估结果
def plot_segmentation(image, prediction, label):
    plt.figure(figsize=(15, 5))

    plt.subplot(1, 3, 1)
    plt.imshow(np.transpose(image.cpu().numpy(), (1, 2, 0)))
    plt.title('Original Image')

    plt.subplot(1, 3, 2)
    plt.imshow(prediction.cpu().numpy(), cmap='gray')
    plt.title('Predicted Segmentation')

    plt.subplot(1, 3, 3)
    plt.imshow(label.cpu().numpy(), cmap='gray')
    plt.title('True Label')

    plt.show()


images, labels = next(iter(val_loader))
image = images[0].to(device)
label = labels[0].to(device)
output = model(image.unsqueeze(0))
_, predicted = torch.max(output, 1)

plot_segmentation(image, predicted.squeeze(), label)


# 在测试集上测试
model.eval()
total_iou = 0.0
total_pixel_acc = 0.0
num_batches = 0

with torch.no_grad():
    pbar = tqdm(test_loader)
    for i, (images, labels) in enumerate(pbar):
        images, labels = images.to(device), labels.to(device)

        outputs = model(images)
        _, predicted = torch.max(outputs, 1)

        intersection = (predicted == labels).sum().item()
        union = (predicted >= 0).sum().item() + (labels >= 0).sum().item() - intersection
        batch_iou = intersection / union if union > 0 else 0

        total_iou += batch_iou
        total_pixel_acc += (predicted == labels).sum().item() / (labels.numel())

        num_batches += 1

mean_iou = total_iou / num_batches
mean_pixel_acc = total_pixel_acc / num_batches

print(f'Test Mean IoU: {mean_iou:.4f}, Test Pixel Accuracy: {mean_pixel_acc:.4f}')